package com.castlzl.random.core.leaf.redis;

import com.castlzl.random.core.leaf.LeafSegmentIdGenerator;
import com.castlzl.random.core.leaf.LeafSegmentService;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.Collections;

/**
 * Redis实现的号段分配服务
 * 使用Redis的INCRBY原子操作实现号段分配
 * @author leiziliang
 */
public class RedisLeafSegmentServiceImpl implements LeafSegmentService {

    private final StringRedisTemplate redisTemplate;
    private final long step;
    private final String keyPrefix;

    private static final String SEGMENT_SCRIPT =
            "local key = KEYS[1]\n" +
            "local step = tonumber(ARGV[1])\n" +
            "local current = redis.call('GET', key)\n" +
            "if not current then\n" +
            "    current = redis.call('GET', key .. ':backup')\n" +
            "    if not current then\n" +
            "        -- 尝试从持久化存储中获取\n" +
            "        current = redis.call('GET', key .. ':persistent')\n" +
            "        if not current then\n" +
            "            current = '0'\n" +
            "        end\n" +
            "    end\n" +
            "end\n" +
            "current = tonumber(current)\n" +
            "local next = current + step\n" +
            "-- 更新当前值和备份\n" +
            "redis.call('SET', key, tostring(next))\n" +
            "redis.call('SET', key .. ':backup', tostring(next))\n" +
            "-- 确保持久化存储\n" +
            "redis.call('SET', key .. ':persistent', tostring(next))\n" +
            "redis.call('PERSIST', key)\n" +
            "redis.call('PERSIST', key .. ':backup')\n" +
            "redis.call('PERSIST', key .. ':persistent')\n" +
            "local min = current + 1\n" +
            "local max = next\n" +
            "return {min, max}";


    private final DefaultRedisScript<java.util.List> script;

    @SuppressWarnings("unchecked")
    public RedisLeafSegmentServiceImpl(StringRedisTemplate redisTemplate, long step, String keyPrefix) {
        this.redisTemplate = redisTemplate;
        this.step = step;
        this.keyPrefix = keyPrefix;

        this.script = new DefaultRedisScript<>();
        this.script.setScriptText(SEGMENT_SCRIPT);
        this.script.setResultType(java.util.List.class);
    }

    @Override
    public LeafSegmentIdGenerator.Segment getNextSegment(String bizTag) throws Exception {
        String key = keyPrefix + bizTag;
        
        java.util.List<Long> result = redisTemplate.execute(
                script,
                Collections.singletonList(key),
                String.valueOf(step)
        );

        if (result == null || result.size() != 2) {
            throw new RuntimeException("Failed to get segment from Redis");
        }

        return new LeafSegmentIdGenerator.Segment(result.get(0), result.get(1));
    }
}